#include "norm.hpp"
#include "gcr.h"
#include <math.h>
#include <iostream>
#include <fstream>
#include <mpi.h>
#include <mkl.h>
#include <mkl_cblas.h>
#include "gptl.h"
//#include "pop_halo_c.hpp"
//#include "matrixpro.hpp"

extern "C" void matrixpro_c(float *a,float *b,float *c,int its, int ite, int jts, int jte, int kts, int kte, int jend);
extern "C" void module_halo_mp_glob_updatehalo_real_3d_(float* h_p, int *,int *); 
extern "C" void svrasr( float*  yl,int m, float*  b, int ni, int nk, int nj);
extern "C" void glob_updatehalo_p_(float *h_p, int *m, int *iter_max,int *its, int *ite, int *kts, int *kte, int *jts, int *jte);
extern "C" void matrixpro_x(float *a,double *b,float *c,int its, int ite, int jts, int jte, int kts, int kte, int jend, int ims, int ime,  int jms, int jme);
static double dot2(const float *u, const float *v, const  int n)
{
    double tmp_sum = 0;
    #pragma omp parallel for  reduction(+:tmp_sum) schedule(runtime)
    for(int i=0; i<n; i++) 
        tmp_sum += u[i] * v[i];
    return tmp_sum;//sqrt(tmp);
}

static double norm2(const float *v, const  int n)
{
    double tmp_sum = 0;
    #pragma omp parallel for  reduction(+:tmp_sum) schedule(runtime)
    for(int i=0; i<n; i++) 
        tmp_sum += v[i] * v[i];
    return tmp_sum;//sqrt(tmp);
}

extern "C" void  psolve_main(float ep, float *a0, float *f0, float *cm_helm,int iter_max, double *x0, 
    int dep, int  jdep, int  ids, int  ide, int  jds, int  jde, int  kds, int  kde, 
    int ims, int  ime, int  jms, int  jme, int  kms, int  kme, 
    int its, int  ite, int  jts, int  jte, int  kts, int  kte, int *mytid){

    #include "defvar.hpp"
GPTLstart ("initialization");
    initialize(h_p,NG1);
    initialize(h_ap,NG);
    initialize(h_r,NG);
    m = 0; d = 0.0;
    int tid=*mytid; int zero=0;
    //para_jte1(h_p[index4b(i,k,j,0)]=x0[index_x(i,k,j)];)
    //glob_updatehalo_p_(h_p, &m, &iter_max, &its, &ite, &kts, &kte, &jts, &jte);
    //para_jte1(x0[index_x(i,k,j)]=h_p[index4b(i,k,j,0)];)
            //=x0[index_x(i,k,j)];)

    matrixpro_x(a0, x0,h_ap,its, ite, jts, jte, kts, kte, jend);
    //gcr_init
    para_jend(h_r[index3(i,k,j)]  = f0[index3(i,k,j)]-h_ap[index3(i,k,j)];)
    para_jte(h_p[index4b(i,k,j,zero)] = h_r[index3(i,k,j)];
            h_r0[index3(i,k,j)] = h_r[index3(i,k,j)];)
    GPTLstop ("initialization");
    double ac=0.0,ac1=0.0;
    //para_jend(ac = ac + h_r0[index3(i,k,j)]*h_ap[index3(i,k,j)];)
    c1[0] = norm2(h_r0, NG);
    MPI_Allreduce(c1,c2,1,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
    ac = c2[0]; 
    float omega=0.7, alpha=0.0;
    if(tid==0) printf("\nline=%d, ac=%.5e",m,ac); 
    double init_err=sqrt(ac);
    //max_iteration=20;
    for(int s=0;s<max_iteration;s++){
        m+=iter_max;
        GPTLstart ("gcr_iteration"); 
        glob_updatehalo_p_(h_p, &zero, &iter_max, &its, &ite, &kts, &kte, &jts, &jte);

        para_jte1(h_bp[index3b(i,k,j)] = h_p[index3b(i,k,j)];)
        GPTLstart ("precondition"); 
        svrasr(h_bp,zero, cm_helm, NX, NY, NZ);
        GPTLstop ("precondition"); 
        glob_updatehalo_p_(h_bp, &zero, &iter_max, &its, &ite, &kts, &kte, &jts, &jte); 
        //Ap=A*p
        GPTLstart ("matrixpro"); 
        matrixpro_c(a0, h_bp,h_ap, its, ite, jts, jte, kts, kte, jend);
        GPTLstop ("matrixpro"); 
        //ac1=ac
        //alpha=ac1/<r0,Ap>
        GPTLstart ("dot"); 
        c1[0] = dot2(h_r0,h_ap,NG);
        MPI_Allreduce(c1,c2,1,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
        GPTLstop ("dot"); 
        ac1=ac; ac=c2[0]; alpha=ac1/ac;
        if(tid==0) printf("\nline=%d, ac=%.5e,alpha=%.5e",m,ac,alpha); 
        //s=r-a*Ap
        #pragma omp parallel for private(j,k,i) schedule(runtime) collapse(2)
        para_jte( h_s[index3b(i,k,j)] = h_r[index3(i,k,j)]-alpha*h_ap[index3(i,k,j)];)
       
        //h_s=U\L\ h_s
        glob_updatehalo_p_(h_s, &zero, &iter_max, &its, &ite, &kts, &kte, &jts, &jte);
        #pragma omp parallel for private(j,k,i) schedule(runtime) collapse(2)
        para_jte1(h_bs[index3b(i,k,j)] = h_s[index3b(i,k,j)];)
     
        GPTLstart ("precondition"); 
        svrasr(h_bs,zero, cm_helm, NX, NY, NZ);
        GPTLstop ("precondition"); 
        glob_updatehalo_p_(h_bs, &zero, &iter_max, &its, &ite, &kts, &kte, &jts, &jte); 
     
        //As = A*s 
        GPTLstart ("matrixpro"); 
        matrixpro_c(a0, h_bs,h_as, its, ite, jts, jte, kts, kte, jend);
        GPTLstop ("matrixpro"); 
        
        //omega=<As,s>/<As,As>
        #pragma omp parallel for private(j,k,i) schedule(runtime) collapse(2)
        para_jte(h_s1[index3(i,k,j)] = h_s[index3b(i,k,j)];)
        c11=0.0; c12=0.0;
        GPTLstart ("dot"); 
        c1[0]=norm2(h_as,NG); c1[1]=dot2(h_as,h_s1, NG);
        MPI_Allreduce(c1,c2,2,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
        GPTLstop ("dot"); 
        omega = c2[1]/c2[0];
        if(tid==0) printf("\nline=%d, c1=%.5e,c2=%.5e,omega=%.5e",m,c2[0],c2[1],omega); 
        //x=x+a*p+w*s;
        #pragma omp parallel for private(j,k,i) schedule(runtime) collapse(2)
        para_jte1(
             x0[index_x(i,k,j)]= x0[index_x(i,k,j)] + alpha * h_bp[index3b(i,k,j)]+ omega*h_bs[index3b(i,k,j)];
        )
        // r=s-w*As
        //if(s%5==0){ 
        if(s==-1){ 
        matrixpro_x(a0, x0,h_ar,its, ite, jts, jte, kts, kte, jend, ims, ime, jms, jme);
        initialize(h_r,NG);
        para_jend(h_r[index3(i,k,j)]  = f0[index3(i,k,j)]- h_ar[index3(i,k,j)];)
        }else{
        #pragma omp parallel for private(j,k,i) schedule(runtime) collapse(2)
        para_jend( h_r[index3(i,k,j)] = h_s[index3b(i,k,j)]-omega*h_as[index3(i,k,j)];) 
        }
        //ac=<r0,r>
        //b=(ac/ac1)*(a/w)
        GPTLstart ("dot"); 
        c1[0] = dot2(h_r0,h_r, NG); c1[1] = norm2(h_r, NG);
        MPI_Allreduce(c1,c2,2,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
        GPTLstop ("dot"); 
        ac = c2[0]; d = c2[1];
        float beta=(ac/ac1)*(alpha/omega);
        // printf("tid=%d,d=%.5e\n",tid,c1[1]); 
        if(tid==0) printf("\nline=%d, ac=%.5e,d=%.5e",m,ac,d); 
        //printr();
        if(tid==0) printf("\nline=%d, beta=%.5e,omega=%.5e\n",m,beta,omega); 
   
        //p=r+b*p-b*w*Ap
        #pragma omp parallel for private(j,k,i) schedule(runtime) collapse(2)
        para_jend(
            h_p[index3b(i,k,j)] = h_r[index3(i,k,j)]+ beta *h_p[index3b(i,k,j)] - beta*omega*h_ap[index3(i,k,j)];
        )
        
 
        GPTLstop ("gcr_iteration");  
        if (d <= ep || s == max_iteration-1 ) { 
            if(tid==0)printf("RES of bicgstab is %E in %d iterations\n", d ,s);
            break;}
        if(tid==0)printf("RES of bicgstab is %E in %d iterations\n", d ,s);
        //return;
    }
    //para_jte1(h_p[index4b(i,k,j,0)]=x0[index_x(i,k,j)];)
    //    glob_updatehalo_p_(h_p, &m, &iter_max, &its, &ite, &kts, &kte, &jts, &jte);
    //matrixpro_c(a0, h_p,h_ap,its, ite, jts, jte, kts, kte, jend);
    matrixpro_x(a0, x0,h_ap,its, ite, jts, jte, kts, kte, jend, ims, ime, jms, jme);
    initialize(h_r,NG);
        #pragma omp parallel for private(j,k,i) schedule(runtime) collapse(2)
    para_jend(h_r[index3(i,k,j)]  = f0[index3(i,k,j)]- h_ap[index3(i,k,j)];)

    c1[0] = norm2(h_r, NG);
    MPI_Allreduce(c1,c2,1,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
    c2[0]=(c2[0]);
    if(tid==5|| tid==0)
        printf("RES of bicgstab is %e \n",c2[0]);
    double relres=sqrt(c2[0]);
    if(tid==0) printf("relres=%.4e,Rate=%.4e\n",relres,exp(log(relres/init_err)/m));


	//delete[] h_p;
    //delete[] h_ap; 
   // delete[] h_ar; delete[] h_r;
}

